import numpy as np
import cv2
import math

# 输入A，B，C 3个关节的坐标，计算AB和BC的夹角，用弧度来表示

def distance(A, B):
    if A is None or B is None:
        return 0
    else:
        return math.sqrt((A[0] - B[0]) ** 2 + (A[1] - B[1]) ** 2) #A[0]代表x坐标，A[1]代表y坐标
 
def Angle(A, B, C):
    """计算三点形成的角度，返回值范围为 0 到 2π"""
    # 计算向量
    vector1 = np.array(B) - np.array(A)
    vector2 = np.array(B) - np.array(C)
    
    # 计算角度
    angle = np.arctan2(vector2[1], vector2[0]) - np.arctan2(vector1[1], vector1[0])
    
    # 将角度转换为 0 到 2π 的范围
    if angle < 0:
        angle += 2 * np.pi
    
    return angle


# 计算两个关节角度的 余弦相似度
def cosine_similarity(angle1, angle2):
    # 计算余弦相似度
    similarity = math.cos(abs(angle1 - angle2))
    return similarity


frameWidth = 640
frameHeight = 640

class PoseKeypoints:
    def __init__(self, results, frame_width, frame_height, person_id):
        """
        初始化关键点类
        keypoints: YOLO模型��出的关键点数据
        frame_width: 画布宽度
        frame_height: 画布高度
        person_id: 指定的人物索引
        """
        self.frame_width = frame_width 
        self.frame_height = frame_height
        self.person_id = person_id
        self.keypoints = results.keypoints[self.person_id].xyn.cpu().numpy() if results.keypoints is not None else None

    def get_keypoint(self, keypoint_idx):
        """获取指定人物的指定关键点坐标"""
        if self.keypoints is None:
            return None
            
        try:
            point = self.keypoints[0][keypoint_idx].tolist()
            x = point[0] * self.frame_width
            y = point[1] * self.frame_height
            return [x, y]
        except:
            print(f"获取关键点错误")
            return None

    def get_nose(self):
        """获取鼻子坐标"""
        return self.get_keypoint(0)
        
    def get_left_eye(self):
        """获取左眼坐标"""
        return self.get_keypoint(1)
        
    def get_right_eye(self):
        """获取右眼坐标"""
        return self.get_keypoint(2)
        
    def get_left_ear(self):
        """获取左耳坐标"""
        return self.get_keypoint(3)
        
    def get_right_ear(self):
        """获取右耳坐标"""
        return self.get_keypoint(4)
        
    def get_left_shoulder(self):
        """获取左肩坐标"""
        return self.get_keypoint(5)
        
    def get_right_shoulder(self):
        """获取右肩坐标"""
        return self.get_keypoint(6)
        
    def get_left_elbow(self):
        """获取左肘坐标"""
        return self.get_keypoint(7)
        
    def get_right_elbow(self):
        """获取右肘坐标"""
        return self.get_keypoint(8)
        
    def get_left_wrist(self):
        """获取左手腕坐标"""
        return self.get_keypoint(9)
        
    def get_right_wrist(self):
        """获取右手腕坐标"""
        return self.get_keypoint(10)
        
    def get_left_hip(self):
        """获取左髋部坐标"""
        return self.get_keypoint(11)
        
    def get_right_hip(self):
        """获取右髋部坐标"""
        return self.get_keypoint(12)
        
    def get_left_knee(self):
        """获取左膝坐标"""
        return self.get_keypoint(13)
        
    def get_right_knee(self):
        """获取右膝坐标"""
        return self.get_keypoint(14)
        
    def get_left_ankle(self):
        """获取左踝坐标"""
        return self.get_keypoint(15)
        
    def get_right_ankle(self):
        """获取右踝坐标"""
        return self.get_keypoint(16)
    
   
    def PoseShow(self, frame, person_idx=0,show_angles=True):
        """在图像上绘制人体骨架和关节角度，如果关键点不存在则跳过该部分绘制"""
        try:
            # 获取所有关键点
            nose = self.get_nose()
            left_eye = self.get_left_eye()
            right_eye = self.get_right_eye()
            left_ear = self.get_left_ear()
            right_ear = self.get_right_ear()
            left_shoulder = self.get_left_shoulder()
            right_shoulder = self.get_right_shoulder()
            left_elbow = self.get_left_elbow()
            right_elbow = self.get_right_elbow()
            left_wrist = self.get_left_wrist()
            right_wrist = self.get_right_wrist()
            left_hip = self.get_left_hip()
            right_hip = self.get_right_hip()
            left_knee = self.get_left_knee()
            right_knee = self.get_right_knee()
            left_ankle = self.get_left_ankle()
            right_ankle = self.get_right_ankle()

            # 绘制骨架连线
            pairs = [
                (left_eye, right_eye), (left_eye, nose), (right_eye, nose),
                (left_eye, left_ear), (right_eye, right_ear),
                (left_shoulder, right_shoulder), (left_shoulder, left_elbow),
                (right_shoulder, right_elbow), (left_elbow, left_wrist),
                (right_elbow, right_wrist), (left_shoulder, left_hip),
                (right_shoulder, right_hip), (left_hip, right_hip),
                (left_hip, left_knee), (right_hip, right_knee),
                (left_knee, left_ankle), (right_knee, right_ankle)
            ]

            # 绘制骨架线条
            for pair in pairs:
                if pair[0] is not None and pair[1] is not None:  # 确保两个端点都存在
                    pt1 = (int(pair[0][0]), int(pair[0][1]))
                    pt2 = (int(pair[1][0]), int(pair[1][1]))
                    # 检查坐标是否为 (0, 0)
                    if pt1 != (0, 0) and pt2 != (0, 0):
                        try:
                            cv2.line(frame, pt1, pt2, (214, 0, 231), 2)
                        except (ValueError, TypeError):
                            continue

            # 绘制关节点
            keypoints = [nose, left_eye, right_eye, left_ear, right_ear,
                        left_shoulder, right_shoulder, left_elbow, right_elbow,
                        left_wrist, right_wrist, left_hip, right_hip,
                        left_knee, right_knee, left_ankle, right_ankle]

            for point in keypoints:
                if point is not None:  # 确保关键点存在
                    pt = (int(point[0]), int(point[1]))
                    # 检查坐标是否为 (0, 0)
                    if pt != (0, 0):
                        try:
                            cv2.circle(frame, pt, 5, (174, 0, 255), -1)
                        except (ValueError, TypeError):
                            continue
            if show_angles:
                # 绘制左右肩部角度
                draw_angle_arc(frame, left_elbow, left_shoulder, right_shoulder)
                draw_angle_arc(frame, right_elbow, right_shoulder, left_shoulder)

                # 绘制左右肘部角度
                draw_angle_arc(frame, left_shoulder, left_elbow, left_wrist)
                draw_angle_arc(frame, right_shoulder, right_elbow, right_wrist)

                # 绘制左右髋部角度
                draw_angle_arc(frame, left_shoulder, left_hip, right_hip)
                draw_angle_arc(frame, right_shoulder, right_hip, left_hip)

                # 绘制左右膝盖相对髋部的角度
                draw_angle_arc(frame, left_knee, left_hip, right_hip)
                draw_angle_arc(frame, right_knee, right_hip, left_hip)

                # 绘制左右膝盖角度
                draw_angle_arc(frame, left_hip, left_knee, left_ankle)
                draw_angle_arc(frame, right_hip, right_knee, right_ankle)


        except Exception as e:
            print(f"绘制过程中发生错误: {str(e)}")
            
        return frame
    
class JointAngles:
    def __init__(self, pose):
        """初始化关节角度类"""
        self.pose = pose
        
    def get_shoulder_angle(self):
        """获取左肘-左肩-右肩的角度"""
        left_elbow = self.pose.get_left_elbow()
        left_shoulder = self.pose.get_left_shoulder() 
        right_shoulder = self.pose.get_right_shoulder()
        return Angle(left_elbow, left_shoulder, right_shoulder) if all((left_elbow, left_shoulder, right_shoulder)) else None

    def get_right_shoulder_angle(self):
        """获取右肘-右肩-左肩的角度"""
        right_elbow = self.pose.get_right_elbow()
        right_shoulder = self.pose.get_right_shoulder()
        left_shoulder = self.pose.get_left_shoulder()
        return Angle(right_elbow, right_shoulder, left_shoulder) if all((right_elbow, right_shoulder, left_shoulder)) else None

    def get_left_elbow_angle(self):
        """获取左肘角度"""
        left_wrist = self.pose.get_left_wrist()
        left_elbow = self.pose.get_left_elbow()
        left_shoulder = self.pose.get_left_shoulder()
        return Angle(left_wrist, left_elbow, left_shoulder) if all((left_wrist, left_elbow, left_shoulder)) else None

    def get_right_elbow_angle(self):
        """获取右肘角度"""
        right_shoulder = self.pose.get_right_shoulder()
        right_elbow = self.pose.get_right_elbow()
        right_wrist = self.pose.get_right_wrist()
        return Angle(right_shoulder, right_elbow, right_wrist) if all((right_shoulder, right_elbow, right_wrist)) else None

    def get_left_hip_angle(self):
        """获取左肩-左髋-右髋的角度"""
        left_shoulder = self.pose.get_left_shoulder()
        left_hip = self.pose.get_left_hip()
        right_hip = self.pose.get_right_hip()
        return Angle(left_shoulder, left_hip, right_hip) if all((left_shoulder, left_hip, right_hip)) else None

    def get_right_hip_angle(self):
        """获取右肩-右髋-左髋的角度"""
        right_shoulder = self.pose.get_right_shoulder()
        right_hip = self.pose.get_right_hip()
        left_hip = self.pose.get_left_hip()
        return Angle(right_shoulder, right_hip, left_hip) if all((right_shoulder, right_hip, left_hip)) else None

    def get_left_knee_hip_angle(self):
        """获取左膝-左髋-右髋的角度"""
        left_knee = self.pose.get_left_knee()
        left_hip = self.pose.get_left_hip()
        right_hip = self.pose.get_right_hip()
        return Angle(left_knee, left_hip, right_hip) if all((left_knee, left_hip, right_hip)) else None

    def get_right_knee_hip_angle(self):
        """获取右膝-右髋-左髋的角度"""
        right_knee = self.pose.get_right_knee()
        right_hip = self.pose.get_right_hip()
        left_hip = self.pose.get_left_hip()
        return Angle(right_knee, right_hip, left_hip) if all((right_knee, right_hip, left_hip)) else None

    def get_left_knee_angle(self):
        """获取左膝角度"""
        left_hip = self.pose.get_left_hip()
        left_knee = self.pose.get_left_knee()
        left_ankle = self.pose.get_left_ankle()
        return Angle(left_hip, left_knee, left_ankle) if all((left_hip, left_knee, left_ankle)) else None

    def get_right_knee_angle(self):
        """获取右膝角度"""
        right_hip = self.pose.get_right_hip()
        right_knee = self.pose.get_right_knee()
        right_ankle = self.pose.get_right_ankle()
        return Angle(right_hip, right_knee, right_ankle) if all((right_hip, right_knee, right_ankle)) else None

    def get_all_angles(self):
        """返回所有关节角度的列表，如果角度不存在则用0代替"""
        angles = []
        
        # 获取所有角度
        angles.append(("左肩角度", self.get_shoulder_angle() or 0))
        angles.append(("右肩角度", self.get_right_shoulder_angle() or 0))
        angles.append(("左肘角度", self.get_left_elbow_angle() or 0))
        angles.append(("右肘角度", self.get_right_elbow_angle() or 0))
        angles.append(("左腰角度", self.get_left_hip_angle() or 0))
        angles.append(("右腰角度", self.get_right_hip_angle() or 0))
        angles.append(("左髋角度", self.get_left_knee_hip_angle() or 0))
        angles.append(("右髋角度", self.get_right_knee_hip_angle() or 0))
        angles.append(("左膝角度", self.get_left_knee_angle() or 0))
        angles.append(("右膝角度", self.get_right_knee_angle() or 0))
        
        return angles

def draw_angle_arc(frame, A, B, C):
    """绘制关节角度的弧线和角度值"""
    # 检查所有点是否存在
    if A is not None and B is not None and C is not None:
        try:
            angle_rad = Angle(A, B, C)  # 使用参数 A, B, C
            
            # 将弧度转换为角度，并确保在 0 到 360 的范围内
            angle_deg = np.degrees(angle_rad)  # 转换为度
            if angle_deg < 0:
                angle_deg += 360
            
            if angle_deg:
                # 计算弧线的起始和结束角度
                start_angle = np.degrees(np.arctan2(B[1] - A[1], B[0] - A[0]))  # 从 A 到 B 的角度
                end_angle = start_angle + angle_deg  # 从 A 到 B 的角度加上计算的角度
                
                # 绘制角度弧线
                center = (int(B[0]), int(B[1]))  # 使用 B 作为肘部
                radius = 30

                # cv2.ellipse(frame, center, (radius, radius), 0, start_angle,end_angle , (119, 119, 255), 2)  # 使用起始和结束角度
                # 显示角度值
                cv2.putText(frame, f"{angle_deg:.1f}", 
                            (center[0]-20, center[1]-20),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (214, 0, 231), 2)
        except (ValueError, TypeError):
            pass
    else:
        return


def compare_poses(frame, pose1, pose2, threshold=0.5):
    """比较两个姿态的相似度并在画面上标记
    
    Args:
        frame: 视频帧
        pose1: 第一个人的PoseKeypoints对象
        pose2: 第二个人的PoseKeypoints对象
        threshold: 相似度阈值,低于此值的关节会被标记为红色
    """
    # 获取两个pose的所有角度
    angles1 = JointAngles(pose1).get_all_angles()
    angles2 = JointAngles(pose2).get_all_angles()
    
    total_similarity = 0
    valid_angles = 0
    
    # 遍历每个关节角度进行比较
    for (name1, angle1), (name2, angle2) in zip(angles1, angles2):
        if angle1 and angle2:  # 如果两个角度都存在
            # 计算余弦相似度
            similarity = cosine_similarity(angle1, angle2)
            total_similarity += similarity
            valid_angles += 1
            
            # 如果相似度低于阈值,标记对应的关节点为红色
            if similarity < threshold:
                # 根据关节名称��取对应的关节点并标记
                if "肩" in name1:
                    if "左" in name1:
                        pt = pose2.get_left_shoulder()
                    else:
                        pt = pose2.get_right_shoulder()
                elif "肘" in name1:
                    if "左" in name1:
                        pt = pose2.get_left_elbow()
                    else:
                        pt = pose2.get_right_elbow()
                elif "腰" in name1:
                    if "左" in name1:
                        pt = pose2.get_left_hip()
                    else:
                        pt = pose2.get_right_hip()
                elif "髋" in name1:
                    if "左" in name1:
                        pt = pose2.get_left_hip()
                    else:
                        pt = pose2.get_right_hip()
                elif "膝" in name1:
                    if "左" in name1:
                        pt = pose2.get_left_knee()
                    else:
                        pt = pose2.get_right_knee()
                
                if pt is not None:
                    cv2.circle(frame, (int(pt[0]), int(pt[1])), 8, (0,0,255), -1)
    
    # 计算平均相似度
    avg_similarity = total_similarity / valid_angles if valid_angles > 0 else 0
    # print(avg_similarity)
    
    # 在第二个人头顶上方显示总体相似度
    head_point = pose2.get_nose()
    if head_point is not None:
        cv2.putText(frame, f"Similarity: {avg_similarity:.2f}",
                    (int(head_point[0]-50), int(head_point[1]-30)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,255,0), 2)
    # cv2.putText(frame, f"Similarity: {avg_similarity:.2f}",
    #             (10,30),
    #             cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,255,0), 2)
    
    return avg_similarity

